#!/usr/bin/env python
# -*- coding:utf-8 -*-
__author__ = 'wshu'
__version__ = '1.0'
"""
    ***********************************
    *  @filename : NessusApi.py
    *  @Author : wshu
    *  @CodeDate : 2020/3/17 15:12
    *  @Software : PyCharm
    ***********************************
"""
import atexit
import urllib3
import requests
import json as serialize

from JQSetting.models import Scanner


class NessusCli(object):
    """
    定制Nessus连接
    """
    def __init__(self, api_akey='', api_skey='', insecure=False):
        self.scan_name = ''
        self.scanner_id = ''
        self.api_akey = None
        self.api_skey = None
        self.use_api = False
        self.res = dict()
        self.scan_template_uuid = ''
        self.scan_uuid = ''

        if insecure and hasattr(requests, 'packages'):
            urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)

        if api_akey and api_skey:
            self.api_akey = api_akey
            self.api_skey = api_skey
            self.use_api = True
        else:
            print('无效的扫描器[Access_Key]和[Secret_Key]')
            atexit.register(self.action, action='session', method='delete')

        # ------------
    def get_scanner_info(self):
        """
        获取HKMSettings中扫描器配置信息
        :return:
        """
        scanner = Scanner.objects.filter(id=self.scanner_id).first()
        self.url = scanner.scanner_url
        self.Access_Key = scanner.scanner_apikey
        self.Secret_Key = scanner.scanner_apisec
        # ------------
        print("[+]扫描器地址: {url}".format(url=self.url))
        print("[+]AccessKey: {ak}".format(ak=self.Access_Key))
        print("[+]SecretKey: {sk}".format(sk=self.Secret_Key))

    # ------------
    def action(self, action, method, payload=None, json_req=True, download=False, private=False):

        headers = {'X-ApiKeys': 'accessKey=' + self.Access_Key +
                                '; secretKey=' + self.Secret_Key}

        assert (isinstance(payload, dict)), 'payload must is dict type'
        if json_req and payload:
            headers.update({'Content-type': 'application/json',
                            'Accept': 'text/plain'})
            payload = serialize.dumps(payload)

        url = "{}/{}".format(self.url, action)
        try:
            rqs = requests.request(method, url, data=payload,
                                   verify=False, headers=headers)
            if not download and rqs.text:
                self.res = rqs.json()
            elif not rqs.text:
                self.res = dict()

            if rqs.status_code != 200:
                print("*****************START ERROR*****************")
                print("HEADERS :")
                print(headers)
                print("URL     : %s " % url)
                print("METHOD  : %s" % method)
                print("RESPONSE: %d" % rqs.status_code)
                print("\n")

            if download:
                return rqs.text
        except requests.exceptions.SSLError as sslError:
            raise Exception('{} for {}'.format(sslError, url))
        except requests.exceptions.ConnectionError:
            raise Exception("Could not connect to {}.\nExiting!\n".format(url))
    # ------------
    def policy_details(self, scanner_id):
        '''
        Retrieves details of an existing policy.
        '''
        self.policy_id = scanner_id
        self.action(action="policies/" + str(self.policy_id), method="GET")
        return dict((p['name'], p['template_uuid']) for p in self.res['policies'])

    # ------------
    def _scan_template_uuid(self, name):
        '''
        Get the template ID. This provides the default settings for the policy.
        '''
        self.action(action="editor/scan/templates", method="GET")
        for template in self.res["templates"]:
            if template["name"] == name:
                self.scan_template_uuid = template["uuid"]
                break

    # ------------
    def scan_add(self, target, template='custom', name='', start=''):
        """
        :param target: 扫描目标
        :param template: 扫描模板
        :param name: 扫描名称
        :param start:
        :return:
        """
        self.scan_template_uuid(name=template)
        self._scan_tag()

        settings = dict()

        t_target = target.replace(',', '\n')
        self.target = target.replace(',', '')

        self.scan_name = name

        scan = dict(uuid=self.scan_template_uuid)

        # Static items- some could be dynamic, but it's overkill
        settings.update({"launch": "ON_DEMAND"})
        settings.update({"description": "Created with REST API"})
        settings.update({"file_targets": ""})
        settings.update({"filters": []})
        settings.update({"emails": ""})
        settings.update({"filter_type": ""})

        settings.update({"scanner_id": str(self.scanner_id)})
        settings.update({"name": self.scan_name})

        if start:
            settings.update({"starttime": start})
            settings.update({"rrules": "FREQ=ONETIME"})

        scan.update(dict(settings=settings))

        self.action(action="scans", method='POST', payload=scan)

        self.scan_uuid = self.res["scan"]["uuid"]
